package gainratio

import (
	"math"
)

// данные одной категории должны быть расположены по строкам
// (первый индекс слайса дает все данные 1 категории)
func CalcGeinRatio(data [][]float64, targetIndex int) []float64 {

	result := make([]float64, 0)
	for i := 0; i < len(data); i++ {
		if i == targetIndex {
			continue
		}
		result = append(result, gainRatio(data[i], data[targetIndex]))
	}
	return result
}

func gainRatio(x []float64, target []float64) float64 {
	data := groupByX(x, target)

	entropyBefore := calcEntropy(target)

	entropyAfther := 0.
	splitInfo := 0.

	for _, ti := range data {
		//нтропия группы
		a := calcEntropy(ti) * (float64(len(ti)) / float64(len(target)))
		entropyAfther += a

		//разбиение информации
		wi := float64(len(ti)) / float64(len(target))
		splitInfo += wi * math.Log2(wi)

	}

	//fmt.Println("split", splitInfo)
	informationGain := entropyBefore - entropyAfther
	//fmt.Println("information gain", informationGain)

	return informationGain / -splitInfo

}

func groupByX(x, target []float64) map[float64][]float64 {
	result := make(map[float64][]float64)

	for i := range x {
		key := x[i]
		result[key] = append(result[key], target[i])
	}

	return result
}

func calcEntropy(class []float64) float64 {
	counts := counter(class)

	totalSamples := len(class)

	entropy := 0.
	for _, count := range counts {
		probability := float64(count) / float64(totalSamples)
		entropy -= probability * math.Log2(probability)
	}
	return entropy
}

func counter(class []float64) map[float64]int {
	res := make(map[float64]int)
	for _, v := range class {
		res[v]++
	}
	return res
}
